"""
核心数据处理与计算逻辑：数据生成、预处理、模型构建与采样、诊断与表格输出。
"""
import os
import time
import gc
import random
import numpy as np
import pandas as pd
import pymc as pm
import arviz as az
from PIL import Image


def _attach_posterior_predictive(trace, posterior_predictive):
	"""Try to merge posterior_predictive InferenceData into trace robustly.

	Prefer arviz.concat; if that fails, try to set the posterior_predictive group
	directly on the trace object.
	"""
	try:
		# preferred: concat InferenceData objects
		trace = az.concat(trace, posterior_predictive)
		return trace
	except Exception:
		try:
			# fallback: try to copy groups directly
			if hasattr(posterior_predictive, "posterior_predictive"):
				trace.posterior_predictive = posterior_predictive.posterior_predictive
			if hasattr(posterior_predictive, "observed_data") and hasattr(posterior_predictive, "observed_data"):
				# rarely present; attempt to attach
				trace.observed_data = getattr(posterior_predictive, "observed_data", None)
		except Exception:
			# give up quietly; caller will handle missing posterior_predictive
			pass
	return trace


def _sample_and_attach(model, config):
	"""Run sampling and posterior predictive in one place and attach result.

	This centralizes how we choose cores/chains and guarantees the returned
	InferenceData contains a posterior_predictive group when possible.
	"""
	# decide cores: prefer n_threads when parallel enabled
	sample_cores = config.get("n_threads") if config.get("parallel") else config.get("cores", 1)
	trace = pm.sample(
		draws=config.get("draws", 500),
		tune=config.get("tune", 1000),
		cores=sample_cores,
		chains=config.get("chains", 1),
		init="jitter+adapt_diag",
		random_seed=config.get("random_seed", 42),
		return_inferencedata=True,
		model=model,
	)
	# posterior predictive
	posterior_predictive = pm.sample_posterior_predictive(
		trace, model=model, return_inferencedata=True
	)
	trace = _attach_posterior_predictive(trace, posterior_predictive)
	return trace


def generate_apc_data(config):
	n_ages = config.get("n_ages", 10)
	n_periods = config.get("n_periods", 5)
	n_cohorts = config.get("n_cohorts", 10)
	n_individuals_per_group = config.get("n_individuals_per_group", 200)

	age_groups = np.arange(20, 20 + n_ages * 5, 5)
	periods = np.arange(2010, 2010 + n_periods * 4, 4)
	data = []
	for period in periods:
		for age in age_groups:
			cohort = period - age
			n_individuals = n_individuals_per_group
			gender = np.random.binomial(1, 0.5, n_individuals)
			region = np.random.choice(["东部", "中部", "西部"], n_individuals)
			education = np.clip(np.random.normal(12, 3, n_individuals), 6, 22)
			age_effect = 0.1 * (age - 40) - 0.005 * (age - 40) ** 2
			period_effect = 0.05 * (period - 2020)
			cohort_effect = 0.03 * (cohort - 1980)
			gender_effect = 0.05 * gender
			region_effect = np.zeros(n_individuals)
			region_effect[region == "东部"] = 0.15
			region_effect[region == "中部"] = 0.05
			education_effect = 0.03 * education
			interaction_effect = 0.001 * (age - 40) * (education - 12)
			base_income = 3000
			income = (
				base_income
				* (1 + age_effect)
				* (1 + period_effect)
				* (1 + cohort_effect)
				* (1 + gender_effect)
				* (1 + region_effect)
				* (1 + education_effect)
				* (1 + interaction_effect)
			)
			income = income + np.random.normal(0, 500, n_individuals)
			group_data = pd.DataFrame({
				"age": [age] * n_individuals,
				"period": [period] * n_individuals,
				"cohort": [cohort] * n_individuals,
				"gender": gender,
				"region": region,
				"education": education,
				"income": income,
			})
			data.append(group_data)
	return pd.concat(data, ignore_index=True)


def preprocess_data(data):
	data["age_group"] = pd.Categorical(data["age"])
	data["period_code"] = pd.Categorical(data["period"]).codes
	# 使用 deciles 分箱可能在少量 unique 值下失败，调用者需保证样本量
	data["cohort_bin"] = pd.qcut(data["cohort"], q=10, labels=False)
	data["cohort_code"] = data["cohort_bin"]
	data["gender_code"] = data["gender"]
	data["region_code"] = pd.Categorical(data["region"]).codes
	return data


def build_and_sample_apc_model(data, config):
	n_ages = len(data["age_group"].cat.categories)
	n_periods = len(data["period_code"].unique())
	n_cohorts = len(data["cohort_code"].unique())
	age_idx = data["age_group"].cat.codes.values
	period_idx = data["period_code"].values
	cohort_idx = data["cohort_code"].values
	y = data["income"].values
	with pm.Model() as apc_model:
		mu = pm.Normal("mu", mu=0, sigma=100)
		# 支持 sum-to-zero 约束（通过 config['sum_to_zero']=True 启用）
		if config.get("sum_to_zero", False):
			age_raw = pm.Normal("age_raw", mu=0, sigma=50, shape=n_ages)
			age_effect = pm.Deterministic("age_effect", age_raw - pm.math.mean(age_raw))
			period_raw = pm.Normal("period_raw", mu=0, sigma=50, shape=n_periods)
			period_effect = pm.Deterministic("period_effect", period_raw - pm.math.mean(period_raw))
			cohort_raw = pm.Normal("cohort_raw", mu=0, sigma=50, shape=n_cohorts)
			cohort_effect = pm.Deterministic("cohort_effect", cohort_raw - pm.math.mean(cohort_raw))
		else:
			age_effect = pm.Normal("age_effect", mu=0, sigma=50, shape=n_ages)
			period_effect = pm.Normal(
				"period_effect", mu=0, sigma=50, shape=n_periods
			)
			cohort_effect = pm.Normal("cohort_effect", mu=0, sigma=50, shape=n_cohorts)
		mu_linear = mu + age_effect[age_idx] + period_effect[period_idx] + cohort_effect[cohort_idx]
		sigma = pm.HalfCauchy("sigma", beta=50)
		y_obs = pm.Normal("y_obs", mu=mu_linear, sigma=sigma, observed=y)
		# centralize sampling + posterior_predictive merging
		trace_apc = _sample_and_attach(apc_model, config)
	return trace_apc, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y


def build_and_sample_full_model(
	data, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y, config
):
	gender = data["gender_code"].values
	region = data["region_code"].values
	education = data["education"].values
	with pm.Model() as full_model:
		mu = pm.Normal("mu", mu=0, sigma=100)
		if config.get("sum_to_zero", False):
			age_raw = pm.Normal("age_raw", mu=0, sigma=50, shape=n_ages)
			age_effect = pm.Deterministic("age_effect", age_raw - pm.math.mean(age_raw))
			period_raw = pm.Normal("period_raw", mu=0, sigma=50, shape=n_periods)
			period_effect = pm.Deterministic("period_effect", period_raw - pm.math.mean(period_raw))
			cohort_raw = pm.Normal("cohort_raw", mu=0, sigma=50, shape=n_cohorts)
			cohort_effect = pm.Deterministic("cohort_effect", cohort_raw - pm.math.mean(cohort_raw))
		else:
			age_effect = pm.Normal("age_effect", mu=0, sigma=50, shape=n_ages)
			period_effect = pm.Normal("period_effect", mu=0, sigma=50, shape=n_periods)
			cohort_effect = pm.Normal("cohort_effect", mu=0, sigma=50, shape=n_cohorts)
		beta_gender = pm.Normal("beta_gender", mu=0, sigma=10)
		beta_region = pm.Normal("beta_region", mu=0, sigma=10)
		beta_education = pm.Normal("beta_education", mu=0, sigma=10)
		mu_linear = (
			mu
			+ age_effect[age_idx]
			+ period_effect[period_idx]
			+ cohort_effect[cohort_idx]
			+ beta_gender * gender
			+ beta_region * region
			+ beta_education * education
		)
		sigma = pm.HalfCauchy("sigma", beta=50)
		y_obs = pm.Normal("y_obs", mu=mu_linear, sigma=sigma, observed=y)
		# centralize sampling + posterior_predictive merging
		trace_full = _sample_and_attach(full_model, config)
	return trace_full


def build_and_sample_full_model_with_interaction(
	data, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y, config
):
	gender = data["gender_code"].values
	region = data["region_code"].values
	education = data["education"].values
	age_edu_interaction = data["age_group"].cat.codes.values * education
	with pm.Model() as full_model:
		mu = pm.Normal("mu", mu=0, sigma=100)
		if config.get("sum_to_zero", False):
			age_raw = pm.Normal("age_raw", mu=0, sigma=50, shape=n_ages)
			age_effect = pm.Deterministic("age_effect", age_raw - pm.math.mean(age_raw))
			period_raw = pm.Normal("period_raw", mu=0, sigma=50, shape=n_periods)
			period_effect = pm.Deterministic("period_effect", period_raw - pm.math.mean(period_raw))
			cohort_raw = pm.Normal("cohort_raw", mu=0, sigma=50, shape=n_cohorts)
			cohort_effect = pm.Deterministic("cohort_effect", cohort_raw - pm.math.mean(cohort_raw))
		else:
			age_effect = pm.Normal("age_effect", mu=0, sigma=50, shape=n_ages)
			period_effect = pm.Normal("period_effect", mu=0, sigma=50, shape=n_periods)
			cohort_effect = pm.Normal("cohort_effect", mu=0, sigma=50, shape=n_cohorts)
		beta_gender = pm.Normal("beta_gender", mu=0, sigma=10)
		beta_region = pm.Normal("beta_region", mu=0, sigma=10)
		beta_education = pm.Normal("beta_education", mu=0, sigma=10)
		beta_age_edu = pm.Normal("beta_age_edu", mu=0, sigma=10)
		mu_linear = (
			mu
			+ age_effect[age_idx]
			+ period_effect[period_idx]
			+ cohort_effect[cohort_idx]
			+ beta_gender * gender
			+ beta_region * region
			+ beta_education * education
			+ beta_age_edu * age_edu_interaction
		)
		sigma = pm.HalfCauchy("sigma", beta=50)
		y_obs = pm.Normal("y_obs", mu=mu_linear, sigma=sigma, observed=y)
		# centralize sampling + posterior_predictive merging
		trace_full = _sample_and_attach(full_model, config)
	return trace_full


def print_descriptive_table(data):
	desc = data[["age", "period", "cohort", "education", "income"]].describe().T
	desc = desc[["count", "mean", "std", "min", "max"]]
	desc.columns = ["样本数", "均值", "标准差", "最小值", "最大值"]
	print("\n表1：主要变量描述性统计")
	print(desc.to_string(float_format="%.2f"))


def print_apc_identification_note():
	print("\n【APC模型识别策略说明】")
	print(
		"APC（年龄-时期-队列）模型存在完全共线性（age + cohort = period），导致三者效应无法唯一分离，称为APC识别问题。"
	)
	print(
		"本脚本采用贝叶斯方法，通过对年龄、时期、队列效应分别施加零均值正态先验（如 N(0, 50)），"
	)
	print(
		"在先验约束下实现参数的可识别性。这样得到的效应为相对效应（相对于整体均值的偏离），"
	)
	print("有助于缓解识别问题并获得稳定的估计结果。")
	print("如需更严格的识别，可进一步采用sum-to-zero约束或设置参考组等方法。\n")


def model_diagnostics(trace, model_name, fig_save_path: str = None):
	print(f"\n【{model_name} 采样诊断】")
	summary = az.summary(trace, round_to=2)
	print(summary[["mean", "sd", "r_hat", "ess_bulk", "ess_tail"]])
	# 后验预测检验（PPC）和残差分析需要 posterior_predictive 数据
	has_ppc = False
	try:
		# InferenceData typically has attribute posterior_predictive
		if hasattr(trace, "posterior_predictive") and trace.posterior_predictive is not None:
			has_ppc = True
	except Exception:
		has_ppc = False

	if not has_ppc:
		print(f"\n{model_name} 未找到 posterior_predictive，已跳过 PPC 图与残差分析（可在建模函数中合并 posterior_predictive）。")
		return

	print(f"\n{model_name} 后验预测检验（PPC）图已保存")
	ppc_fig = az.plot_ppc(trace, figsize=(8, 4))
	import matplotlib.pyplot as plt
	plt.tight_layout()
	# 保存 Figure
	try:
		fig = ppc_fig.figure
	except Exception:
		fig = plt.gcf()
	from src.plot import save_fig
	save_fig(fig, f"{model_name}_ppc.png", fig_save_path)

	try:
		if "y_obs" in trace.posterior_predictive:
			y_obs = trace.observed_data["y_obs"].values
			y_pred = trace.posterior_predictive["y_obs"].values.mean(axis=(0, 1))
			plt.figure(figsize=(6, 4))
			plt.hist(residuals := (y_obs - y_pred), bins=30, alpha=0.7)
			plt.title(f"{model_name} 残差分布")
			plt.xlabel("残差")
			plt.ylabel("频数")
			plt.tight_layout()
			save_fig(plt.gcf(), f"{model_name}_residuals.png", fig_save_path)
			print(f"{model_name} 残差分布图已保存")
		else:
			print(f"{model_name} 未找到后验预测数据，无法进行残差分析。")
	except Exception as e:
		print(f"在生成残差分析时出错：{e}")


def print_full_model_diagnosis(trace):
	"""Print a concise diagnostic summary for the full model to the terminal.

	This provides simple checks (R-hat, ESS) and textual guidance similar to
	the project's script template.
	"""
	summary = az.summary(trace, round_to=2)
	rhat_max = summary["r_hat"].max()
	ess_min = summary["ess_bulk"].min()
	print("\n【完整模型诊断结论】")
	if rhat_max < 1.05:
		print("1. 采样收敛性良好（所有参数R-hat < 1.05）。")
	else:
		print("1. 警告：部分参数R-hat >= 1.05，存在收敛风险。")
	if ess_min > 200:
		print("2. 有效样本数充足（最小bulk ESS > 200），参数估计可靠。")
	else:
		print("2. 警告：部分参数有效样本数较低，建议增加采样步数。")
	print("3. 后验预测检验（PPC）和残差分布图已保存，可进一步判断模型拟合优度和假设合理性。")
	print("4. 建议结合PPC图和残差分布，综合评估模型拟合效果。")
	print("\n【完整模型拟合诊断总结】")
	print("1. PPC图显示模型预测分布与观测分布整体较为接近，拟合优度较好。")
	print("2. 残差分布以0为中心，但右侧有长尾，提示对极端高收入的拟合略有不足。")
	print("3. 建议关注极端值的处理，或尝试更灵活的误差分布以提升模型表现。\n")


def save_param_table(trace, param_names):
	summary = az.summary(trace, var_names=param_names, hdi_prob=0.95, round_to=2)
	summary = summary[["mean", "hdi_2.5%", "hdi_97.5%"]]
	summary.columns = ["均值", "95%下界", "95%上界"]
	print("\n【完整模型主要参数估计表】")
	print(summary.to_string(float_format="%.2f"))


def save_inferencedata(trace, name: str, out_dir: str = "outputs"):
	"""Save InferenceData to netcdf under out_dir/name.nc"""
	os.makedirs(out_dir, exist_ok=True)
	path = os.path.join(out_dir, f"{name}.nc")
	try:
		trace.to_netcdf(path)
		print(f"InferenceData 已保存: {path}")
	except Exception as e:
		print(f"保存 InferenceData 失败: {e}")


def model_compare(traces: dict, method: str = "loo"):
	"""Compare models using arviz. Accepts dict{name: InferenceData}.

	Returns an arviz comparison DataFrame.
	"""
	idata_list = list(traces.values())
	names = list(traces.keys())
	try:
		if method == "loo":
			loo_list = [az.loo(idata) for idata in idata_list]
			cmp = az.compare(dict(zip(names, idata_list)), method="stacking")
		else:
			cmp = az.compare(dict(zip(names, idata_list)))
		print("模型比较结果：")
		print(cmp)
		return cmp
	except Exception as e:
		print(f"模型比较失败: {e}")
		return None


def group_ppc_plot(trace, data: pd.DataFrame, group_col: str, var_name: str = "y_obs", fig_save_path: str = None, plot_cfg: dict = None):
	"""Generate PPC per group and save figures. group_col should exist in data."""
	groups = data[group_col].unique()
	for g in groups:
		sub_idx = data[data[group_col] == g].index
		try:
			pp = trace.posterior_predictive[var_name].sel({"chain": slice(None), "draw": slice(None)}).values
			# pp shape: (chain, draw, obs)
			pp_sub = pp.reshape(-1, pp.shape[-1])[:, sub_idx]
			import matplotlib.pyplot as plt
			plt.figure(figsize=(6, 4))
			plt.hist(pp_sub.mean(axis=0), bins=30, alpha=0.7)
			plt.title(f"PPC ({group_col}={g})")
			plt.tight_layout()
			from src.plot import save_fig
			save_fig(plt.gcf(), f"ppc_{group_col}_{g}.png", fig_save_path, plot_cfg=plot_cfg)
			print(f"PPC 图已保存：{group_col}={g}")
		except Exception as e:
			print(f"为组 {g} 生成 PPC 失败: {e}")


def prior_sensitivity_run(data, base_config: dict, prior_stds: list = [10, 25, 50], model_builder=None, var_of_interest: str = "beta_education"):
	"""Run several fits with varying prior std and return a summary table for var_of_interest.

	model_builder: function(data, config) -> trace
	"""
	results = {}
	for s in prior_stds:
		cfg = base_config.copy()
		cfg["prior_std"] = s
		# caller's model_builder should read cfg['prior_std'] to set priors
		trace = model_builder(data, cfg)
		try:
			summary = az.summary(trace, var_names=[var_of_interest], hdi_prob=0.95)
			results[s] = summary[["mean", "hdi_2.5%", "hdi_97.5%"]]
		except Exception as e:
			results[s] = f"error: {e}"
	return results


def print_beta_age_edu(trace):
	summary = az.summary(trace, var_names=["beta_age_edu"], hdi_prob=0.95, round_to=4)
	mean = summary.loc["beta_age_edu", "mean"]
	hdi_low = summary.loc["beta_age_edu", "hdi_2.5%"]
	hdi_high = summary.loc["beta_age_edu", "hdi_97.5%"]
	print("\n【年龄×教育交互项估计】")
	print(f"beta_age_edu 均值：{mean:.4f}，95%区间：[{hdi_low:.4f}, {hdi_high:.4f}]")
	if hdi_low > 0 or hdi_high < 0:
		print("结论：交互项显著，说明不同年龄组教育对收入的边际效应存在差异。")
	else:
		print("结论：交互项不显著，不同年龄组教育对收入的边际效应无明显差异。")


def gender_group_analysis(data, config, font, fig_save_path: str = None):
	traces = {}
	n_ages = len(data['age_group'].cat.categories)
	age_categories = data['age_group'].cat.categories
	for gender_code, gender_label in zip([0, 1], ["女性", "男性"]):
		print(f"\n【分组回归：{gender_label}】")
		sub_data = data[data['gender_code'] == gender_code].copy()
		sub_data = preprocess_data(sub_data)
		age_idx = sub_data['age_group'].cat.codes.values
		period_idx = sub_data['period_code'].values
		cohort_idx = sub_data['cohort_code'].values
		y = sub_data['income'].values
		trace, *_ = build_and_sample_apc_model(sub_data, config)
		traces[gender_label] = trace
		save_param_table(trace, ['mu', 'age_effect', 'period_effect', 'cohort_effect'])
		from src.plot import plot_effect
		plot_effect(
			trace,
			'age_effect',
			n_ages,
			age_categories,
			f'{gender_label}年龄效应估计 (95% HDI)',
			f"age_effect_{gender_label}.png",
			font,
			fig_save_path,
		)
	from src.plot import compare_effects
	compare_effects(
		traces['女性'], traces['男性'], 'age_effect', n_ages, age_categories,
		'不同性别组年龄效应对比', 'compare_age_effect_gender.png', font, fig_save_path, label1='女性', label2='男性'
	)
